import pandas as pd
import argparse
import os

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='将细胞信息（lineage和disease）合并到目标CSV文件中',
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument('--input', '-i', 
                       required=True,
                       help='要写入细胞信息的CSV文件路径')
    parser.add_argument('--info', '-f', 
                       required=True,
                       help='细胞信息CSV文件路径')
    parser.add_argument('--output', '-o',
                       help='输出文件路径（可选，默认覆盖输入文件）')
    parser.add_argument('--no-drop-duplicates', '-n', 
                       action='store_true',
                       help='不删除输出文件中的重复细胞系行（默认会删除）')
    
    return parser.parse_args()

def normalize_column_name(col_name):
    """规范化列名：去除空格、转为小写"""
    return col_name.strip().lower()

def find_matching_column(columns, target_name):
    """在列名列表中查找匹配的列名（忽略大小写和空格）"""
    normalized_columns = {normalize_column_name(col): col for col in columns}
    normalized_target = normalize_column_name(target_name)
    return normalized_columns.get(normalized_target)

def remove_duplicate_cell_lines(df, cell_line_column):
    """
    检查并删除DataFrame中重复的细胞系行
    
    参数:
        df: 要处理的DataFrame
        cell_line_column: 细胞系列名
    
    返回:
        处理后的DataFrame
    """
    # 检查是否存在重复的细胞系
    if df[cell_line_column].duplicated().any():
        # 获取重复的细胞系列表
        duplicate_cells = df[df[cell_line_column].duplicated()][cell_line_column].tolist()
        print(f"发现 {len(duplicate_cells)} 个重复的细胞系")
        print(f"重复的细胞系: {duplicate_cells}")
        
        # 删除重复项，保留第一个出现的
        original_rows = len(df)
        df_clean = df.drop_duplicates(subset=[cell_line_column], keep='first')
        removed_rows = original_rows - len(df_clean)
        print(f"已删除 {removed_rows} 行重复数据，处理后包含 {len(df_clean)} 个唯一细胞系")
        
        return df_clean
    else:
        print("未发现重复的细胞系")
        return df

def merge_cell_info(input_file, info_file, output_file=None, drop_duplicates=True):
    """
    将细胞信息文件中的lineage和disease列合并到目标文件中
    
    参数:
        input_file: 要写入细胞信息的CSV文件路径
        info_file: 细胞信息CSV文件路径
        output_file: 输出文件路径，None表示覆盖输入文件
        drop_duplicates: 是否删除重复的细胞系行（默认True）
    """
    # 如果未指定输出文件，默认覆盖输入文件
    if output_file is None:
        output_file = input_file
    
    # 验证文件是否存在
    if not os.path.exists(input_file):
        print(f"错误：输入文件不存在: {input_file}")
        return False
    
    if not os.path.exists(info_file):
        print(f"错误：细胞信息文件不存在: {info_file}")
        return False
    
    try:
        # 读取细胞信息文件
        cell_info_df = pd.read_csv(info_file)
        print(f"细胞信息文件读取完成，包含 {len(cell_info_df)} 行数据")
        print(f"细胞信息文件原始列名: {list(cell_info_df.columns)}")

        # 读取目标CSV文件
        target_df = pd.read_csv(input_file)
        print(f"目标文件读取完成，包含 {len(target_df)} 行数据")
        print(f"目标文件原始列名: {list(target_df.columns)}")

        # 规范化列名
        cell_info_df.columns = [normalize_column_name(col) for col in cell_info_df.columns]
        target_df.columns = [normalize_column_name(col) for col in target_df.columns]
        
        print(f"规范化后细胞信息文件列名: {list(cell_info_df.columns)}")
        print(f"规范化后目标文件列名: {list(target_df.columns)}")
        
        # 检查是否包含必要的列
        cell_line_col_info = find_matching_column(cell_info_df.columns, 'cell_line')
        if not cell_line_col_info:
            print(f"错误：细胞信息文件中缺少 'cell_line' 列（现有列: {list(cell_info_df.columns)}）")
            return False
        
        lineage_col_info = find_matching_column(cell_info_df.columns, 'lineage')
        disease_col_info = find_matching_column(cell_info_df.columns, 'disease')
        if not lineage_col_info or not disease_col_info:
            print(f"错误：细胞信息文件中缺少 'lineage' 或 'disease' 列（现有列: {list(cell_info_df.columns)}）")
            return False
        
        cell_line_col_target = find_matching_column(target_df.columns, 'cell_line')
        if not cell_line_col_target:
            print(f"错误：目标文件中缺少 'cell_line' 列（现有列: {list(target_df.columns)}）")
            return False
        
        # 检查目标文件是否已经包含lineage和disease列
        lineage_col_target = find_matching_column(target_df.columns, 'lineage')
        disease_col_target = find_matching_column(target_df.columns, 'disease')
        
        if lineage_col_target and disease_col_target:
            print(f"注意：目标文件已包含 'lineage' 和 'disease' 列，将使用细胞信息文件中的数据覆盖")
            # 创建临时DataFrame用于合并
            temp_info_df = cell_info_df[[cell_line_col_info, lineage_col_info, disease_col_info]].copy()
            
            # 重命名列以匹配目标文件
            if lineage_col_info != lineage_col_target:
                temp_info_df = temp_info_df.rename(columns={lineage_col_info: lineage_col_target})
            if disease_col_info != disease_col_target:
                temp_info_df = temp_info_df.rename(columns={disease_col_info: disease_col_target})
            
            # 重命名cell_line列以匹配目标文件
            if cell_line_col_info != cell_line_col_target:
                temp_info_df = temp_info_df.rename(columns={cell_line_col_info: cell_line_col_target})
                merge_key = cell_line_col_target
            else:
                merge_key = cell_line_col_info
            
            # 删除目标文件中现有的lineage和disease列
            target_df = target_df.drop(columns=[lineage_col_target, disease_col_target])
            
            # 执行合并
            merged_df = pd.merge(target_df, temp_info_df, on=merge_key, how='left')
            
            # 验证合并后的列名
            print(f"合并后的文件列名: {list(merged_df.columns)}")
            
            # 检查是否有细胞系在细胞信息文件中找不到匹配
            if lineage_col_target in merged_df.columns:
                missing_cells = merged_df[merged_df[lineage_col_target].isna()][merge_key].tolist()
                if missing_cells:
                    print(f"警告：以下 {len(missing_cells)} 个细胞系在细胞信息文件中找不到匹配:")
                    print(missing_cells)
                else:
                    print("所有细胞系都在细胞信息文件中找到匹配")
        else:
            # 如果列名不同，重命名以确保匹配
            if cell_line_col_info != cell_line_col_target:
                print(f"注意：将目标文件中的 '{cell_line_col_target}' 列重命名为 '{cell_line_col_info}' 以进行合并")
                target_df = target_df.rename(columns={cell_line_col_target: cell_line_col_info})
                merge_key = cell_line_col_info
            else:
                merge_key = cell_line_col_info
            
            # 检查是否存在重复的细胞系
            if cell_info_df[merge_key].duplicated().any():
                print("警告：细胞信息文件中存在重复的细胞系")
                # 显示重复的细胞系
                duplicate_cells = cell_info_df[cell_info_df[merge_key].duplicated()][merge_key].tolist()
                print(f"重复的细胞系: {duplicate_cells}")
                # 保留第一个出现的细胞系信息
                cell_info_df = cell_info_df.drop_duplicates(subset=[merge_key], keep='first')
                print(f"已移除重复项，保留 {len(cell_info_df)} 个唯一细胞系")

            # 执行合并操作
            merged_df = pd.merge(target_df, cell_info_df[[merge_key, lineage_col_info, disease_col_info]], 
                                on=merge_key, how='left')

            # 检查是否有细胞系在细胞信息文件中找不到匹配
            if lineage_col_info in merged_df.columns:
                missing_cells = merged_df[merged_df[lineage_col_info].isna()][merge_key].tolist()
                if missing_cells:
                    print(f"警告：以下 {len(missing_cells)} 个细胞系在细胞信息文件中找不到匹配:")
                    print(missing_cells)
                else:
                    print("所有细胞系都在细胞信息文件中找到匹配")
        
        # 确定最终的lineage和disease列名
        final_lineage_col = lineage_col_target if lineage_col_target else lineage_col_info
        final_disease_col = disease_col_target if disease_col_target else disease_col_info
        
        # 重新排列列的顺序，将lineage和disease放在前面
        if merge_key in merged_df.columns and final_lineage_col in merged_df.columns and final_disease_col in merged_df.columns:
            columns_order = [merge_key, final_lineage_col, final_disease_col] + \
                          [col for col in merged_df.columns if col not in [merge_key, final_lineage_col, final_disease_col]]
            merged_df = merged_df[columns_order]
        else:
            print("警告：无法重新排列列顺序，某些必要的列可能丢失")
            print(f"合并后的可用列: {list(merged_df.columns)}")
            
        # 如果启用了删除重复项功能，处理合并后的DataFrame
        if drop_duplicates:
            print("\n开始检查并删除重复的细胞系行...")
            merged_df = remove_duplicate_cell_lines(merged_df, merge_key)

        # 保存合并后的文件
        merged_df.to_csv(output_file, index=False)
        print(f"合并后的文件已保存至: {output_file}")
        print(f"合并后的文件包含 {len(merged_df)} 行数据")

        # 显示合并结果的前几行进行验证
        print("\n合并结果预览:")
        print(merged_df.head())
        
        return True
        
    except Exception as e:
        import traceback
        print(f"处理过程中出错: {str(e)}")
        print("详细错误信息:")
        traceback.print_exc()
        return False

if __name__ == "__main__":
    # 解析命令行参数
    args = parse_arguments()
    
    # 执行合并操作 - 默认删除重复项，除非指定了--no-drop-duplicates
    success = merge_cell_info(args.input, args.info, args.output, not args.no_drop_duplicates)
    
    # 根据执行结果设置退出码
    import sys
    sys.exit(0 if success else 1)